# Copyright 2020 - 2022 MONAI Consortium
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
import torch.nn as nn

from .swin_uneter import SwinTransformer as SwinViT
from monai.utils import ensure_tuple_rep
from monai import transforms
from einops import rearrange

transform = transforms.Compose(
    [
        transforms.RandAdjustContrast(prob=0.7, gamma=[0.7, 1.5]),
        transforms.RandShiftIntensity(0.4, prob=0.4),
        transforms.RandGaussianNoise(prob=0.5, mean=0.05, std=0.05),
    ]
)


def aug_rand(samples):
    samples = transform(samples)
    return samples


class SSLHead(nn.Module):
    def __init__(self, args, upsample="vae", dim=768):
        super(SSLHead, self).__init__()
        patch_size = ensure_tuple_rep(2, args.spatial_dims)
        window_size = ensure_tuple_rep(7, args.spatial_dims)

        self.swinViT = SwinViT(
            in_chans=args.in_channels,
            embed_dim=args.feature_size,
            window_size=window_size,
            patch_size=patch_size,
            depths=[2, 2, 2, 2],
            num_heads=[3, 6, 12, 24],
            mlp_ratio=4.0,
            qkv_bias=True,
            drop_rate=0.0,
            attn_drop_rate=0.0,
            drop_path_rate=args.dropout_path_rate,
            norm_layer=torch.nn.LayerNorm,
            use_checkpoint=args.use_checkpoint,
            spatial_dims=args.spatial_dims,
            use_dynamic_mask_ratio=args.dynamic_masking,
            use_hierarchical_mask=args.hierarchical_masking,
            basic_mask_ratio=args.basic_mask_ratio,
            scale=args.scale,
            drop_ratio=args.drop_ratio
        )

        self.conv4 = nn.Sequential(
            nn.Conv3d(dim, dim // 2, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 2),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(dim // 2, dim // 4, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 4),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(dim // 4, dim // 8, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 8),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(dim // 8, dim // 16, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 16),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(dim // 16, dim // 16, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 16),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(dim // 16, args.in_channels, kernel_size=1, stride=1),
        )

        self.conv3 = nn.Sequential(
            nn.Conv3d(384, 192, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(192),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(192, 96, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(96),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(96, 48, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(48),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(48, 24, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(24),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(24, args.in_channels, kernel_size=1, stride=1),
        )

        self.conv2 = nn.Sequential(
            nn.Conv3d(192, 96, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 2),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(96, 48, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 4),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(48, 24, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(24),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(24, args.in_channels, kernel_size=1, stride=1),
        )

        self.conv1 = nn.Sequential(
            nn.Conv3d(96, 48, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(dim // 2),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(48, 24, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(24),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(24, args.in_channels, kernel_size=1, stride=1),
        )

        self.conv0 = nn.Sequential(
            nn.Conv3d(48, 24, kernel_size=3, stride=1, padding=1),
            nn.InstanceNorm3d(24),
            nn.LeakyReLU(),
            nn.Upsample(scale_factor=2, mode="trilinear", align_corners=False),
            nn.Conv3d(24, args.in_channels, kernel_size=1, stride=1),
        )

    def copy_weight(self):
        for (_, ema_param), (_, model_param) in zip(self.swinViT_contrastive.named_parameters(), self.swinViT.named_parameters()):
            ema_param.data = model_param.data

    def forward(self, x):
        x0_out, x1_out, x2_out, x3_out, x4_out = self.swinViT(x.contiguous())

        _, c, h, w, d = x0_out.shape

        x_rec0 = x0_out.flatten(start_dim=2, end_dim=4)
        x_rec0 = x_rec0.view(-1, c, h, w, d)
        x_rec0 = self.conv0(x_rec0)

        _, c, h, w, d = x1_out.shape

        x_rec1 = x1_out.flatten(start_dim=2, end_dim=4)
        x_rec1 = x_rec1.view(-1, c, h, w, d)
        x_rec1 = self.conv1(x_rec1)

        _, c, h, w, d = x2_out.shape

        x_rec2 = x2_out.flatten(start_dim=2, end_dim=4)
        x_rec2 = x_rec2.view(-1, c, h, w, d)
        x_rec2 = self.conv2(x_rec2)

        _, c, h, w, d = x3_out.shape

        x_rec3 = x3_out.flatten(start_dim=2, end_dim=4)
        x_rec3 = x_rec3.view(-1, c, h, w, d)
        x_rec3 = self.conv3(x_rec3)

        _, c, h, w, d = x4_out.shape

        x_rec4 = x4_out.flatten(start_dim=2, end_dim=4)
        x_rec4 = x_rec4.view(-1, c, h, w, d)
        x_rec4 = self.conv4(x_rec4)
        return x_rec0, x_rec1, x_rec2, x_rec3, x_rec4

